#include "pathbuilder.h"

using namespace std;
using namespace cv;

int PathBuilder::GetRegions(cv::Mat &image, std::vector<uint> &labels)
{
	uint label = 0;
	for (int i = 0; i < image.rows; ++i)
	{
		for (int j = 0; j < image.cols; ++j)
		{
			if (labels[i * image.cols + j] == 0 && image.at<uchar>(i, j) == 255)
			{
				label++;
				PutLabel(image, labels, label, i, j);
			}
		}
	}
	return label;
}

void PathBuilder::PutLabel(cv::Mat &image, std::vector<uint> &labels, uint label, uint i, uint j)
{
	if (labels[i * image.cols + j] == 0 && image.at<uchar>(i, j) == 255)
	{
		labels[i * image.cols + j] = label;

		if (i > 0)
			PutLabel(image, labels, label, i - 1, j);

		if (i < image.rows - 1)
			PutLabel(image, labels, label, i + 1, j);

		if (j > 0)
			PutLabel(image, labels, label, i, j - 1);

		if (j < image.cols - 1)
			PutLabel(image, labels, label, i, j + 1);
	}
}

cv::Mat PathBuilder::VizualizePath(cv::Mat &image, std::vector<uint> &labels, uint objNum)
{
	Mat labeledIm;
	image.copyTo(labeledIm);

	uint fraction = 255.0f / objNum;

	for (int i = 0; i < labeledIm.rows; ++i)
	{
		for (int j = 0; j < labeledIm.cols; ++j)
		{
			if (labels[i * labeledIm.cols + j] != 0)
			{
				labeledIm.at<uchar>(i, j) = labels[i * labeledIm.cols + j] * fraction;
			}
		}
	}
	return labeledIm;
}

std::map<uint, std::pair<uint, uint>> PathBuilder::GetCenter(std::vector<uint> & labels, uint objNum, Size size)
{
	map<uint, pair<uint, uint>> objCenters;

	for (int i = 0; i < objNum; ++i)
	{
		objCenters.emplace(i, pair<uint, uint>(0, 0));
	}

	for (int i = 0; i < size.height; ++i)
	{
		for (int j = 0; j < size.width; ++j)
		{
			if (labels[i * size.width + j] != 0)
			{
				objCenters[labels[i * size.width + j] - 1].first += i;
				objCenters[labels[i * size.width + j] - 1].second += j;
			}
		}
	}

	map<uint, uint> area = GetArea(labels, objNum, size);

	for (int k = 0; k < objNum; ++k)
	{
		objCenters[k].first /= area[k];
		objCenters[k].second /= area[k];
	}

	return objCenters;
}

std::map<uint, uint> PathBuilder::GetArea(std::vector<uint> & labels, uint objNum, Size size)
{
	map<uint, uint> area;

	for (int i = 0; i < objNum; ++i)
	{
		area.emplace(i, 0);
	}

	for (int i = 0; i < size.height; ++i)
	{
		for (int j = 0; j < size.width; ++j)
		{
			if (labels[i * size.width + j] != 0)
				area[labels[i * size.width + j] - 1]++;
		}
	}
	return area;
}

std::vector<std::pair<int, int>> PathBuilder::FindNearestNeighbour(std::map<uint, std::pair<uint, uint>> centerCoord1,
	std::map<uint, std::pair<uint, uint>> centerCoord2, map<uint, uint> area1)
{
	uint objNum1 = centerCoord1.size();
	uint objNum2 = centerCoord2.size();

	vector<float> distance(objNum1 * objNum2, 0);

	auto metric = [](pair<uint, uint> coord1, pair<uint, uint> coord2) -> float {
		int dx = (coord1.first - coord2.first) * (coord1.first - coord2.first);
		int dy = (coord1.second - coord2.second) * (coord1.second - coord2.second);
		return (sqrt(dx + dy));
	};

	auto cit = centerCoord1.begin();
	for (int i = 0; i < objNum1; ++i, ++cit)
	{
		for (int j = 0; j < objNum2; ++j)
		{
			distance[i * objNum2 + j] = metric(cit->second, centerCoord2[j]);
		}
	}

	vector<pair<int, int>> neighbours;

	auto ait = area1.begin();
	for (int i = 0; i < objNum1; ++i, ++ait)
	{
		pair<float, int> min(FLT_MAX, -1);

		for (int j = 0; j < objNum2; ++j)
		{
			if (distance[i * objNum2 + j] < min.first)
			{
				min.first = distance[i * objNum2 + j];
				min.second = j;
			}
		}

		if (min.second == -1)
		{
			neighbours.emplace_back(ait->first, -1);
			continue;
		}

		float area = ait->second;
		float radius = sqrt(area / M_PI);

		if (min.first < 2 * radius)
			neighbours.emplace_back(ait->first, min.second);
		else
			neighbours.emplace_back(ait->first, -1);
	}



	return neighbours;
}

cv::Mat PathBuilder::DrawPath(cv::Mat & image, std::vector<std::pair<int, int>> neighbours, 
	std::map<uint, std::pair<uint, uint>> coord1, std::map<uint, std::pair<uint, uint>> coord2)
{
	for (int k = 0; k < neighbours.size(); ++k)
	{
		int prev = neighbours[k].first;
		int next = neighbours[k].second;

		if (next == -1)
		{
			trajectory.erase(prev);
			continue;
		}

		pair<uint, uint> xy1 = coord1[prev];
		pair<uint, uint> xy2 = coord2[next];

		auto tit = trajectory.find(prev);
		if (tit != trajectory.end())
		{
			DrawLine(tit->second, xy1, xy2);
		}
		else
		{
			vector<pair<uint, uint>> path;
			DrawLine(path, xy1, xy2);
			trajectory.emplace(prev, path);
		}
	}

	for (auto tit = trajectory.begin(); tit != trajectory.end(); ++tit)
	{
		for (int i = 0; i < tit->second.size(); ++i)
		{
			auto path = tit->second;

			if (image.channels() == 1)
			{
				image.at<uchar>(path[i].first, path[i].second) = 128;
			}
			else if (image.channels() == 3)
			{
				Vec3b value = image.at<Vec3b>(path[i].first, path[i].second);
				value.val[2] = 255;
				image.at<Vec3b>(path[i].first, path[i].second) = value;
			}
		}
	}

	return image;
}

void PathBuilder::DrawLine(vector<pair<uint, uint>> & path, pair<uint, uint> coord1, pair<uint, uint> coord2)
{
	// Bresenham's line algorithm
	int x1 = coord1.first;
	int y1 = coord1.second;

	int x2 = coord2.first;
	int y2 = coord2.second;

	int deltax = abs(x2 - x1);
	int deltay = abs(y2 - y1);

	bool flag = deltay > deltax;

	if (flag)
	{
		swap(x1, y1);
		swap(x2, y2);
	}

	if (x1 > x2)
	{
		swap(x1, x2);
		swap(y1, y2);
	}

	float error = 0;
	float deltaerr = deltay;

	int sign = -1;
	if (y1 < y2)
	{
		sign = 1;
	}
	
	int y = y1;

	for (int x = x1; x < x2; ++x)
	{
		if (flag)
		{
			path.emplace_back(y, x);
		}
		else
		{
			path.emplace_back(x, y);
		}
		
		error += deltaerr;
		if (2 * error >= deltax)
		{
			y += sign;
			error -= deltax;
		}
	}
}


void PathBuilder::ChangeLabels(std::vector<std::pair<int, int>> & neighbours,
	std::map<uint, std::pair<uint, uint>> & centerCoord, std::map<uint, uint> & area)
{
	map<uint, pair<uint, uint>> newCoord;
	map<uint, uint> newArea;

	for (int i = 0; i < neighbours.size(); ++i)
	{
		int prev = neighbours[i].first;
		int next = neighbours[i].second;

		if (prev == next)
		{
			newCoord.emplace(prev, centerCoord[prev]);
			newArea.emplace(prev, area[prev]);
			centerCoord.erase(prev);
			area.erase(prev);
			continue;
		}
		else if (next != -1)
		{
			newCoord.emplace(prev, centerCoord[next]);
			newArea.emplace(prev, area[next]);
			centerCoord.erase(next);
			area.erase(next);
			neighbours[i].second = neighbours[i].first;
		}
	}

	if (!centerCoord.empty())
	{
		auto ait = area.begin();
		auto cit = centerCoord.begin();

		for (; ait != area.end(), cit != centerCoord.end(); ++ait, ++cit)
		{
			newArea.emplace(++objectAmount, ait->second);
			newCoord.emplace(objectAmount, cit->second);
		}
	}

	area.clear();
	area = newArea;
	centerCoord.clear();
	centerCoord = newCoord;
}


cv::Mat PathBuilder::BuildPath(std::string path, uint amount)
{
	Mat im1 = imread(path + "16d.bmp", CV_LOAD_IMAGE_GRAYSCALE);

	vector<uint> labels(im1.cols * im1.rows, 0);
	int objNum = GetRegions(im1, labels);
	objectAmount = objNum;

	//	Mat labeled = VizualizePath(im1, labels, objNum1);

	map<uint, pair<uint, uint>> centerCoord1 = GetCenter(labels, objNum, im1.size());
	map<uint, uint> area1 = GetArea(labels, objNum, im1.size());

	// 	for (int k = 0; k < objNum1; ++k)
	// 	{
	// 		labeled.at<uchar>(cm[k].first, cm[k].second) = 0;
	// 	}

	cv::Size s;
	s.height = 240;
	s.width = 320;
	string file = "16-36.avi";
	VideoWriter wr = VideoWriter(file, CV_FOURCC_PROMPT, 15.0, s);
	wr << im1;


	vector<pair<int, int>> neighbours;
	Mat res;
	for (int i = 1; i <= amount; ++i)
	{
		//read next frame
		Mat im2 = imread(path + to_string(i) + "d.bmp", CV_LOAD_IMAGE_GRAYSCALE);

		//clear old data
		labels.clear();
		labels = vector<uint>(im2.rows * im2.cols, 0);

		//labeling the frame
		objNum = GetRegions(im2, labels);
		map<uint, pair<uint, uint>> centerCoord2 = GetCenter(labels, objNum, im2.size());
		map<uint, uint> area2 = GetArea(labels, objNum, im2.size());

		//find neighbours
		neighbours = FindNearestNeighbour(centerCoord1, centerCoord2, area1);
		ChangeLabels(neighbours, centerCoord2, area2);

		//draw trajectory
		res = DrawPath(im2, neighbours, centerCoord1, centerCoord2);

		//make a video
		wr << im2;

		//
		centerCoord1.clear();
		centerCoord1 = centerCoord2;
		centerCoord2.clear();

		area1.clear();
		area1 = area2;
		area2.clear();
	}
	return im1;
}


int PathBuilder::GetSynteticDataForWFA()
{
	string truth = "data/dataset1/Heads/Hdep-";
	string data = "data/dataset1/Dep/dep-";
	string result = "data/dataset1/NewDep/dep-";

	for (int i = 0; i <= 2383; ++i)
	{
		Mat im_data = imread(data + to_string(i) + ".png", CV_LOAD_IMAGE_GRAYSCALE);
		Mat im_truth = imread(truth + to_string(i) + ".bmp", CV_LOAD_IMAGE_GRAYSCALE);

		vector<uint> labels(im_truth.cols * im_truth.rows, 0);
		int objNum = GetRegions(im_truth, labels);

		if (objNum == 0)
		{
			cv::imwrite(result + to_string(i) + ".bmp", im_data);
			continue;
		}

		map<uint, vector<pair<int, int>>> pixels;
		
		for (int j = 1; j <= objNum; ++j)
		{
			pixels.emplace(j, vector<pair<int, int>>());
		}

		for (int k = 0; k < im_truth.rows; ++k)
		{
			for (int n = 0; n < im_truth.cols; ++n)
			{
				if (labels[k * im_truth.cols + n] != 0)
				{
					pixels[labels[k * im_truth.cols + n]].emplace_back(k, n);
				}
			}
		}

		map<uint, int> minimum;
		for (int k = 1; k <= objNum; ++k)
		{
			int tmp = INT_MAX;
			int medium = 0;

			for (int q = 0; q < pixels[k].size(); ++q)
			{
				pair<int, int> coord = pixels[k][q];
				if (im_data.at<uchar>(coord.first, coord.second) < tmp)
				{
					tmp = im_data.at<uchar>(coord.first, coord.second);
				}
				medium += im_data.at<uchar>(coord.first, coord.second);
			}

			medium /= pixels[k].size();
			minimum.emplace(k, tmp - 1);

			tmp -= 1;

			for (int q = 0; q < pixels[k].size(); ++q)
			{
				pair<int, int> coord = pixels[k][q];
				im_data.at<uchar>(coord.first, coord.second) = medium - 1;
			}
		}

		cv::imwrite(result + to_string(i) + ".bmp", im_data);

		labels.clear();
		pixels.clear();
	}
	
	return 0;
}



int getSynteticDataForPath()
{

	for (int k = 0; k < 1500; ++k)
	{
		string name = "data/dataset1/GroundTruth/Hdep-" + to_string(k) + ".bmp";

		cv::Mat mat = cv::imread(name, CV_LOAD_IMAGE_COLOR);

		for (int i = 0; i < mat.rows; ++i)
		{
			for (int j = 0; j < mat.cols; ++j)
			{
				Vec3b intensity = mat.at<Vec3b>(i, j);
				uchar blue = intensity.val[0];
				uchar green = intensity.val[1];
				uchar red = intensity.val[2];

				if (red == 255 && (red != green && red != blue))
				{
					intensity.val[0] = 255;
					intensity.val[1] = 255;
					mat.at<Vec3b>(i, j) = intensity;
				}
				else
				{
					intensity.val[0] = 0;
					intensity.val[1] = 0;
					intensity.val[2] = 0;
					mat.at<Vec3b>(i, j) = intensity;
				}
			}
		}

		string res = "data/dataset2/Heads/Hdep-" + to_string(k) + ".bmp";
		imwrite(res, mat);
	}

	return 0;
} 